#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2022/12/29 18:24
# @Author : HuangBenHao

from py2neo import Graph, NodeMatcher
from pprint import pprint
import json


def fix_pro(data):
    for item in dict(data):
        value = data[item].strip()
        value = value.replace(' ', '').replace('\r', '\\r').replace('\n', '\\n')
        data[item] = value


def create_json(n, r, e):
    result = {"p": {}}
    start = {}
    end = {}
    segment = {}

    start['identity'] = n.identity
    start['labels'] = [str(n.labels).replace(':', '')]
    start['properties'] = {}
    for item in dict(n):
        start['properties'][item] = n[item]

    end['identity'] = e.identity
    end['labels'] = [str(e.labels).replace(':', '')]
    end['properties'] = {}
    for item in dict(e):
        end['properties'][item] = e[item]

    segment['relationship'] = {}
    segment['relationship']['identity'] = r.identity
    segment['relationship']['start'] = start['identity']
    segment['relationship']['end'] = end['identity']
    segment['relationship']['type'] = e['name']
    segment['start'] = start
    segment['end'] = end
    segment['relationship']['properties'] = {}
    for item in dict(r):
        segment['relationship']['properties'][item] = r[item]

    result['p']['start'] = start
    result['p']['end'] = end
    result['p']['segments'] = [segment]

    return result


g = Graph("http://localhost:7474", password="bigben123")

data = g.run("match (n:disease)-[r]-(e) where n.name='糖尿病' return n,r,e limit 50").data()

result = []
for item in data:
    single_r = create_json(item['n'], item['r'], item['e'])
    result.append(single_r)

with open(r'./test.json', 'w') as f:
    json.dump(result, f)
